/*
 *
 * Copyright 2012 Mathias Olsson mathias@kompetensum.com
 *
 * This file is dual licensed as either GPL version 2 or Apache License 2.0 at your choice
 * http://www.gnu.org/licenses/old-licenses/gpl-2.0.html
 * http://www.apache.org/licenses/
 *
 * Note that PolarSSL uses GPL with a FOSS License Exception
 *
 */


#include "pbkdf2.h"

#ifndef min
	#define min( a, b ) ( ((a) < (b)) ? (a) : (b) )
#endif

int PKCS5_PBKDF2_HMAC(md_context_t *ctx, const unsigned char *password, size_t plen, const unsigned char *salt, size_t slen,
			const unsigned long iteration_count, const unsigned long key_length, unsigned char *output)
{
	int res = 0;
	
	// Size of the generated digest
	unsigned char md_size = md_get_size(ctx->md_info);
	
	unsigned char* md1 = malloc(md_size);
	if (md1 == NULL) 
	{
		res = -1;
		goto exit;
	}
	
	unsigned char* work = malloc(md_size);
	if (work == NULL) 
	{
		res = -1;
		goto exit;
	}
	
	memset(work,0,md_size);
	
	unsigned long counter = 1;
	unsigned long generated_key_length = 0;
	while (generated_key_length < key_length)
	{
		// U1 ends up in md1 and work
		unsigned char c[4];
		c[0] = (counter >> 24) & 0xff;
		c[1] = (counter >> 16) & 0xff;
		c[2] = (counter >>  8) & 0xff;
		c[3] = (counter >>  0) & 0xff;
		
		res = md_hmac_starts(ctx,password,plen);
		if (res != 0) goto exit;

		res = md_hmac_update(ctx,salt,slen);
		if (res != 0) goto exit;

		res = md_hmac_update(ctx,c,4);
		if (res != 0) goto exit;

		res = md_hmac_finish(ctx,md1);
		if (res != 0) goto exit;
		
		memcpy(work,md1,md_size);
	
		unsigned long ic = 1;
		for (ic = 1; ic < iteration_count; ic++)
		{
			// U2 ends up in md1
			res = md_hmac_starts(ctx,password,plen);
			if (res != 0) goto exit;

			res = md_hmac_update(ctx,md1,md_size);
			if (res != 0) goto exit;

			res = md_hmac_finish(ctx,md1);
			if (res != 0) goto exit;

			// U1 xor U2
			unsigned long i = 0;
			for (i = 0; i < md_size; i++)
			{
				work[i] ^= md1[i];
			}
			
			// and so on until Uiteration_count
		}

		// Copy the generated bytes to the key
		unsigned long bytes_to_write = min((key_length - generated_key_length),md_size);
		memcpy(output + generated_key_length, work, bytes_to_write);
		generated_key_length += bytes_to_write;

		++counter;
	}
	
exit:
		free(md1);
		free(work);
		
		return res;
}

//#define _RUN_TESTS

#ifdef _RUN_TESTS

#include <stdio.h>

typedef struct
{
	char* t;
	char* p;
	int plen;
	char* s;
	int slen;
	int c;
	int dkLen;
	char dk[1024]; // Remember to set this to max dkLen
} testvector;

int do_test(testvector* tv)
{
	printf("Started %s\n",tv->t);
	fflush(stdout);

	const md_info_t* infosha1 = md_info_from_type(POLARSSL_MD_SHA1);
	md_context_t sha1ctx = {0};
	int res = md_init_ctx(&sha1ctx, infosha1);
	if (res != 0)
	{
		return -1;
	}

	char* key = malloc(tv->dkLen);
	if (key == 0)
	{
		return -1;
	}

	res = PKCS5_PBKDF2_HMAC(&sha1ctx,tv->p,tv->plen,tv->s,tv->slen,tv->c,tv->dkLen,key);
	if (res != 0)
	{
		free(key);
		return -1;
	}

	if (memcmp(tv->dk,key,tv->dkLen) != 0)
	{
		// Failed
		free(key);
		return -1;
	}

	free(key);
	return 0;
}


int main(int argc, char** argv)
{
	// Test vectors from RFC 6070

	testvector* tv = 0;
	int res = 0;

/*
    Input:
       P = "password" (8 octets)
       S = "salt" (4 octets)
       c = 1
       dkLen = 20

     Output:
       DK = 0c 60 c8 0f 96 1f 0e 71
            f3 a9 b5 24 af 60 12 06
            2f e0 37 a6             (20 octets)

*/
	testvector t1 = {
			"Test 1",
			"password", 8 , "salt", 4, 1, 20,
			.dk = {0x0c, 0x60, 0xc8, 0x0f, 0x96, 0x1f, 0x0e, 0x71,
			 0xf3, 0xa9, 0xb5, 0x24, 0xaf, 0x60, 0x12, 0x06,
			 0x2f, 0xe0, 0x37, 0xa6}
	};

	tv = &t1;
	res = do_test(tv);
	if (res != 0)
	{
		printf("%s failed\n",tv->t);
		return res;
	}

/*
       Input:
             P = "password" (8 octets)
             S = "salt" (4 octets)
             c = 2
             dkLen = 20

           Output:
             DK = ea 6c 01 4d c7 2d 6f 8c
                  cd 1e d9 2a ce 1d 41 f0
                  d8 de 89 57             (20 octets)

*/

	testvector t2 = {
			"Test 2",
			"password", 8, "salt", 4, 2, 20,
			{0xea, 0x6c, 0x01, 0x4d, 0xc7, 0x2d, 0x6f, 0x8c,
			 0xcd, 0x1e, 0xd9, 0x2a, 0xce, 0x1d, 0x41, 0xf0,
			 0xd8, 0xde, 0x89, 0x57
			}
	};

	tv = &t2;
	res = do_test(tv);
	if (res != 0)
	{
		printf("%s failed\n",tv->t);
		return res;
	}

/*
             Input:
                  P = "password" (8 octets)
                  S = "salt" (4 octets)
                  c = 4096
                  dkLen = 20

                Output:
                  DK = 4b 00 79 01 b7 65 48 9a
                       be ad 49 d9 26 f7 21 d0
                       65 a4 29 c1             (20 octets)


*/
	testvector t3 = {
			"Test 3",
			"password", 8, "salt", 4, 4096, 20,
			{0x4b, 0x00, 0x79, 0x01, 0xb7, 0x65, 0x48, 0x9a,
			 0xbe, 0xad, 0x49, 0xd9, 0x26, 0xf7, 0x21, 0xd0,
			 0x65, 0xa4, 0x29, 0xc1
			}
	};

	tv = &t3;
	res = do_test(tv);
	if (res != 0)
	{
		printf("%s failed\n",tv->t);
		return res;
	}

/*
                  Input:
                     P = "password" (8 octets)
                     S = "salt" (4 octets)
                     c = 16777216
                     dkLen = 20

                   Output:
                     DK = ee fe 3d 61 cd 4d a4 e4
                          e9 94 5b 3d 6b a2 15 8c
                          26 34 e9 84             (20 octets)

*/
	testvector t4 = {
			"Test 4",
			"password", 8, "salt", 4, 16777216, 20,
			{0xee, 0xfe, 0x3d, 0x61, 0xcd, 0x4d, 0xa4, 0xe4,
 			 0xe9, 0x94, 0x5b, 0x3d, 0x6b, 0xa2, 0x15, 0x8c,
			 0x26, 0x34, 0xe9, 0x84
			}
	};

	tv = &t4;
	res = do_test(tv);
	if (res != 0)
	{
		printf("%s failed\n",tv->t);
		return res;
	}

/*
                     Input:
                        P = "passwordPASSWORDpassword" (24 octets)
                        S = "saltSALTsaltSALTsaltSALTsaltSALTsalt" (36 octets)
                        c = 4096
                        dkLen = 25

                      Output:
                        DK = 3d 2e ec 4f e4 1c 84 9b
                             80 c8 d8 36 62 c0 e4 4a
                             8b 29 1a 96 4c f2 f0 70
                             38                      (25 octets)

*/
	testvector t5 = {
			"Test 5",
			"passwordPASSWORDpassword", 24, "saltSALTsaltSALTsaltSALTsaltSALTsalt", 36, 4096, 25,
			{0x3d, 0x2e, 0xec, 0x4f, 0xe4, 0x1c, 0x84, 0x9b,
			 0x80, 0xc8, 0xd8, 0x36, 0x62, 0xc0, 0xe4, 0x4a,
			 0x8b, 0x29, 0x1a, 0x96, 0x4c, 0xf2, 0xf0, 0x70,
			 0x38
			}
	};

	tv = &t5;
	res = do_test(tv);
	if (res != 0)
	{
		printf("%s failed\n",tv->t);
		return res;
	}

/*
                        Input:
                           P = "pass\0word" (9 octets)
                           S = "sa\0lt" (5 octets)
                           c = 4096
                           dkLen = 16

                         Output:
                           DK = 56 fa 6a a7 55 48 09 9d
                                cc 37 d7 f0 34 25 e0 c3 (16 octets)
*/
	testvector t6 = {
			"Test 6",
			"pass\0word", 9, "sa\0lt", 5, 4096, 16,
			{0x56, 0xfa, 0x6a, 0xa7, 0x55, 0x48, 0x09, 0x9d,
			 0xcc, 0x37, 0xd7, 0xf0, 0x34, 0x25, 0xe0, 0xc3,
			}
	};

	tv = &t6;
	res = do_test(tv);
	if (res != 0)
	{
		printf("%s failed\n",tv->t);
		return res;
	}

	printf("All tests successful\n");
	return 0;
}

#endif

